from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Signature import pkcs1_15
from Crypto.Hash import SHA256
import base64


def rsa_generate_keys():
    key = RSA.generate(2048)
    private_key = key.export_key()
    public_key = key.publickey().export_key()
    return private_key, public_key

def rsa_encrypt(text, public_key):
    recipient_key = RSA.import_key(public_key)
    cipher_rsa = PKCS1_OAEP.new(recipient_key)
    text_bytes = text.encode('utf-8')
    encrypted_bytes = cipher_rsa.encrypt(text_bytes)
    encrypted_text = base64.b64encode(encrypted_bytes).decode('utf-8')
    return encrypted_text

def rsa_decrypt(text, private_key):
    key = RSA.import_key(private_key)
    cipher_rsa = PKCS1_OAEP.new(key)
    encrypted_bytes = base64.b64decode(text)
    decrypted_bytes = cipher_rsa.decrypt(encrypted_bytes)
    decrypted_text = decrypted_bytes.decode('utf-8')
    return decrypted_text
def rsa_sign(text, private_key):
    """生成文本的RSA数字签名"""
    key = RSA.import_key(private_key)
    h = SHA256.new(text.encode('utf-8'))
    signature = pkcs1_15.new(key).sign(h)
    return base64.b64encode(signature).decode('utf-8')

def rsa_verify(text, signature, public_key):
    """验证文本的RSA数字签名"""
    try:
        key = RSA.import_key(public_key)
        h = SHA256.new(text.encode('utf-8'))
        signature_bytes = base64.b64decode(signature)
        pkcs1_15.new(key).verify(h, signature_bytes)
        return "签名验证成功"
    except (ValueError, TypeError):
        return "签名验证失败"